import torch
from QfUtil.MyDatasets.ViAuDataset.datasetclass import ViAuDataset
from . import dataset_file_path
from .a1_file_version_control_system import result_vcs


def get_loaders(use_gpu=True, version=None):
    """

    :param use_gpu: always use gpu in current version
    :param version: if load version specified, load using version control system
    :return: train_loader, test_loader
    """
    # if load version specified, load using version control system
    if version is not None:
        train_loader = result_vcs.load('train_loader.loader', version=version)
        test_loader = result_vcs.load('test_loader.loader', version=version)
        return train_loader, test_loader

    # load dataset, and init loaders from dataset.
    dataset = ViAuDataset(label='dataset', path=dataset_file_path)
    train_loader, test_loader = dataset.cross_validation_loader(k=2, use_gpu=use_gpu)

    # save result using vcs
    result_vcs.save(train_loader, 'train_loader.loader')
    result_vcs.save(test_loader, 'test_loader.loader')

    return train_loader, test_loader


def get_samples(version=None, scale_rate=1):
    name = 'samples.data'
    if version is not None:
        return result_vcs.load(name, version)

    dataset = ViAuDataset(label='dataset', path=dataset_file_path).cuda()

    vdin = dataset.data[0][:2000][::scale_rate]
    adin = dataset.data[1][:2000][::scale_rate]
    label = dataset.data[2][:2000][::scale_rate]

    samples = (vdin, adin, label)

    result_vcs.save(samples, name, version)

    return samples


def prepare_sample_with_noise(sample_version=None, noise_strength=0, version=None, scale_rate=1):
    name = 'noise_' + str(noise_strength) + '.samples'
    if version is not None:
        return result_vcs.load(name, version)

    vdin, adin, label = get_samples(version=sample_version, scale_rate=scale_rate)

    vdin = torch.randn_like(vdin) * noise_strength * 0.1 + vdin
    adin = torch.randn_like(adin) * noise_strength * 0.1 + adin

    samples = (vdin, adin, label)

    result_vcs.save(samples, name)

    return samples


def prepare_samples_with_series_noise_strength(version=None, sample_version=None, scale_rate=1):
    samples = []
    for noise_strength in range(11):
        samples.append(
            prepare_sample_with_noise(
                sample_version=sample_version,
                noise_strength=noise_strength,
                version=version,
                scale_rate=scale_rate
            )
        )
    return samples


def yield_samples_with_series_noise_strength(version=None, sample_version=None, scale_rate=1):
    """

    :param version:
    :param sample_version:
    :param scale_rate:
    :return: noise_strength, sample
    """
    for noise_strength in range(11):
        sample = prepare_sample_with_noise(
            sample_version=sample_version,
            noise_strength=noise_strength,
            version=version,
            scale_rate=scale_rate
        )
        yield noise_strength, sample


def add_noise_to_(sample, noise_strength, visual=True):
    """

    :param sample:
    :param noise_strength:
    :param visual:
    :return:
    """
    index = 1
    if visual:
        index = 0
    sample[index] = sample[index] + torch.randn_like(sample[index]) * noise_strength
    return sample
